package jwt

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"io/ioutil"
	"os"
	"sync"

	"project/study-go/config"
	"project/study-go/logger"
	"project/study-go/util"

	"github.com/dgrijalva/jwt-go"
)

var (
	JwtMan *JWTManager
)

//JWTManager 封装了 jwt 的 rsa 密钥对
type JWTManager struct {
	keys *rsa.PrivateKey
	mux  sync.Mutex
}

//Init 为了手动控制初始化
func Init() error {
	var (
		keys *rsa.PrivateKey
		err  error
	)
	keyFile := config.Vip.GetString("rsa.key")
	pubKeyFile := config.Vip.GetString("rsa.pubkey")
	//如果不存在rsa密钥对文件，则生成密钥对文件
	if util.FileNotExist(keyFile) || util.FileNotExist(pubKeyFile) {
		//删除旧的密钥文件
		if util.FileExist(keyFile) {
			os.Remove(keyFile)
		} else if util.FileExist(pubKeyFile) {
			os.Remove(pubKeyFile)
		}
		//生成 rsa keys
		keys, err = rsa.GenerateKey(rand.Reader, 2048)
		if err != nil {
			logger.RoutineLogger.Error("generator rsa keys fail.")
			return err
		}
		//将生成的 priv key 以 pem 格式放入文件中
		privateKey := x509.MarshalPKCS1PrivateKey(keys)
		block := &pem.Block{
			Type:    "RSA PRIVATE KEY",
			Headers: nil,
			Bytes:   privateKey,
		}
		//pem格式在内存中
		pemPriKey := pem.EncodeToMemory(block)
		key, err := os.OpenFile(keyFile, os.O_RDWR|os.O_CREATE, 0644)
		if err != nil {
			logger.RoutineLogger.Errorf("create %s fail.", keyFile)
			return err
		}
		defer key.Close()
		key.Write(pemPriKey)
		//公钥遇上步骤一致
		pubkey := x509.MarshalPKCS1PublicKey(&keys.PublicKey)
		block = &pem.Block{
			Type:    "RSA PUBLIC KEY",
			Headers: nil,
			Bytes:   pubkey,
		}
		pemPubKey := pem.EncodeToMemory(block)
		pubKey, err := os.OpenFile(pubKeyFile, os.O_RDWR|os.O_CREATE, 0644)
		if err != nil {
			logger.RoutineLogger.Errorf("create %s fail.", pubKeyFile)
			return err
		}
		defer pubKey.Close()
		pubKey.Write(pemPubKey)
		logger.RoutineLogger.Infof("Generate rsa keys success.")

	} else { //从文件读取 rsa 密钥
		key, err := os.OpenFile(keyFile, os.O_RDONLY, 0444)
		if err != nil {
			logger.RoutineLogger.Errorf("read %s fail.", keyFile)
			return err
		}
		defer key.Close()
		pemPriKey, err := ioutil.ReadAll(key)
		if err != nil {
			logger.RoutineLogger.Errorf("Read pem rsa keys from %s fail.")
		}
		//将 pem 格式转换为 pem.block 格式
		block, _ := pem.Decode(pemPriKey)
		if block == nil {
			logger.RoutineLogger.Errorf("Convert pem rsa keys to pem.block Ext fail.")
			return errors.New("Convert pem rsa keys to pem.block Ext fail.")
		}
		//再使用 pkcs1 解析rsa 密钥
		keys, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			logger.RoutineLogger.Errorf("Parse rsa privatekey from %s fail.", keyFile)
			return err
		}
		logger.RoutineLogger.Infof("Parse rsa keys for jwtmanager success.")
	}

	logger.RoutineLogger.Infof("Generate jwtManager success.")
	JwtMan = &JWTManager{
		keys: keys,
		mux:  sync.Mutex{},
	}
	return nil
}

//Generate 生成 jwt,使用 RS512 来签名
func (j *JWTManager) Generate(claims jwt.Claims) (string, error) {
	j.mux.Lock()
	defer j.mux.Unlock()
	token := jwt.NewWithClaims(jwt.SigningMethodRS512, claims)
	return token.SignedString(j.keys)
}

//Parse 对 jwt 进行解析和校验
func (j *JWTManager) Parse(tokenString string) (*jwt.Token, error) {
	j.mux.Lock()
	defer j.mux.Unlock()
	token, err := jwt.Parse(tokenString, func(t *jwt.Token) (interface{}, error) {
		return &j.keys.PublicKey, nil
	})
	if err != nil {
		return nil, err
	}
	return token, nil
}
